import torch
import keras
import os

from absl import app
from absl import flags

from torch import nn

from keras import layers as L
from keras import backend as K

from transModel import UltraNet_pytorch
from transModel import UltraNet_keras

FLAGS = flags.FLAGS
flags.DEFINE_string('input_model', None, 'Path to the input model.')
flags.DEFINE_string('output_model', None, 'Path where the converted model will be stored.')

flags.mark_flag_as_required('input_model')
flags.mark_flag_as_required('output_model')

os.environ['CUDA_VISIBLE_DEVICES']='5'

def main(args):
    net_pytorch = UltraNet_pytorch()
    dict = torch.load(FLAGS.input_model)

    net_pytorch.load_state_dict(dict["model"])
    net_pytorch.eval()
    weights_from_torch = net_pytorch.state_dict()

    for k,v in weights_from_torch.items():
        if len(v.shape) == 4:
            weights_from_torch[k] = v.data.numpy().transpose(2, 3, 1, 0)

    net_keras = UltraNet_keras()

    for layer in net_keras.layers:
        current_layer_name = layer.name
        if current_layer_name=='input_layer':
            weights = [weights_from_torch['layers.0.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization':
            weights = [weights_from_torch['layers.1.weight'],weights_from_torch['layers.1.bias'],weights_from_torch['layers.1.running_mean'],weights_from_torch['layers.1.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='conv2d':
            weights = [weights_from_torch['layers.4.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization_1':
            weights = [weights_from_torch['layers.5.weight'],weights_from_torch['layers.5.bias'],weights_from_torch['layers.5.running_mean'],weights_from_torch['layers.5.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='conv2d_1':
            weights = [weights_from_torch['layers.8.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization_2':
            weights = [weights_from_torch['layers.9.weight'],weights_from_torch['layers.9.bias'],weights_from_torch['layers.9.running_mean'],weights_from_torch['layers.9.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='conv2d_2':
            weights = [weights_from_torch['layers.12.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization_3':
            weights = [weights_from_torch['layers.13.weight'],weights_from_torch['layers.13.bias'],weights_from_torch['layers.13.running_mean'],weights_from_torch['layers.13.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='conv2d_3':
            weights = [weights_from_torch['layers.16.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization_4':
            weights = [weights_from_torch['layers.17.weight'],weights_from_torch['layers.17.bias'],weights_from_torch['layers.17.running_mean'],weights_from_torch['layers.17.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='conv2d_4':
            weights = [weights_from_torch['layers.19.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization_5':
            weights = [weights_from_torch['layers.20.weight'],weights_from_torch['layers.20.bias'],weights_from_torch['layers.20.running_mean'],weights_from_torch['layers.20.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='conv2d_5':
            weights = [weights_from_torch['layers.22.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization_6':
            weights = [weights_from_torch['layers.23.weight'],weights_from_torch['layers.23.bias'],weights_from_torch['layers.23.running_mean'],weights_from_torch['layers.23.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='conv2d_6':
            weights = [weights_from_torch['layers.25.weight']]
            layer.set_weights(weights)
        elif current_layer_name=='batch_normalization_7':
            weights = [weights_from_torch['layers.26.weight'],weights_from_torch['layers.26.bias'],weights_from_torch['layers.26.running_mean'],weights_from_torch['layers.26.running_var']]
            layer.set_weights(weights)

        elif current_layer_name=='output_layer':
            weights = [weights_from_torch['layers.28.weight'],weights_from_torch['layers.28.bias']]
            layer.set_weights(weights)

        else :
            if not (isinstance(layer,L.Activation) or isinstance(layer,L.MaxPooling2D)):
                raise ValueError('need to initialize this layer, but not founded in pytorch model')
    net_keras.save(FLAGS.output_model)

if __name__ == "__main__":
    app.run(main)